import sys
sys.path.append("/home/fengmq/code/xiaoche_pi_new/smartcar/ilqr_control_pybind/build")
sys.path.append("/home/fengmq/code/xiaoche_pi_new/smartcar/ilqr_control_pybind/build/proto")
import pybind_test
print("success import pybind")
import ilqr_config_pb2
import ilqr_data_pb2
import control_common_pb2
import ilqr_controller_config_pb2
print("success import proto")
import numpy as np
from google.protobuf import text_format
import math

class State():
    def __init__(self) -> None:
        self.x=None
        self.y=None
        self.theta=None
        self.v=None
    
    def update(self, x ,y ,h, v):
        self.x=x
        self.y=y
        self.theta=h
        self.v=v

class IlqrController():
    def __init__(self) -> None:
        self.roi_time=0.0
        self.lf=1.0
        self.lr=1.0
        self.thr=0.0
        self.brake=0.0
        self.steer=0.0
        self.ego_state=State()
        self.dt=0.1
        self.traj=[]
        self.world_x_list=[]            #只用做坐标系转换
        self.world_y_list=[]            #只用做坐标系转换
        self.world_theta_list=[]        #只用做坐标系转换
        self.traj_base_ego=[]
        self.KPV=1.0
        self.KIV=0.02
        self.KDV=0.0
        self.KPS=1.0
        self.KIS=0.02
        self.KDS=0.0
        self.pre_ds=0.0
        self.pre_dv=0.0
        self.Integral_record_dv=[]
        self.Integral_record_ds=[]
        self.KIntegralMaxSize=5
        self.dv_control_ratio=0.75
        self.KSteerLimit=0.3

    def update_trajectory(self, x, y, h, v, t):
        x_ori_list=x.tolist()
        y_ori_list=y.tolist()
        h_ori_list=h.tolist()
        v_ori_list=v.tolist()
        t_ori_list=t.tolist()
        for i in range(len(h_ori_list)):
            h_ori_list[i]=self.clockwise_0_360_to_pi_negative_pi(h_ori_list[i])
        s_ori_list=[]
        s_now=0.0
        for i in range(len(x_ori_list)):
            if i==0:
                s_ori_list.append(s_now)
            else:
                s_now+=math.sqrt((x_ori_list[i]-x_ori_list[i-1])**2+(y_ori_list[i]-y_ori_list[i-1])**2)
                s_ori_list.append(s_now)

        i=0
        while i<len(x_ori_list):
            if i==0:
                self.traj.append([x_ori_list[i],y_ori_list[i],h_ori_list[i],0.0,v_ori_list[i],0.0])
                self.world_x_list.append(x_ori_list[i])
                self.world_y_list.append(y_ori_list[i])
                self.world_theta_list.append(h_ori_list[i])
                i+=1
            else:
                while i<len(x_ori_list) and t_ori_list[i]>=len(self.traj)*self.dt:
                    tar_time=len(self.traj)*self.dt
                    ratio=(tar_time-t_ori_list[i-1])/(t_ori_list[i]-t_ori_list[i-1])
                    x_temp=x_ori_list[i-1]+ratio*(x_ori_list[i]-x_ori_list[i-1])
                    y_temp=y_ori_list[i-1]+ratio*(y_ori_list[i]-y_ori_list[i-1])
                    h_temp=self.normalize_angle(h_ori_list[i-1]+ratio*(self.normalize_angle(h_ori_list[i]-h_ori_list[i-1])))
                    v_temp=v_ori_list[i-1]+ratio*(v_ori_list[i]-v_ori_list[i-1])
                    s_temp=s_ori_list[i-1]+ratio*(s_ori_list[i]-s_ori_list[i-1])
                    self.traj.append([x_temp,y_temp,h_temp,s_temp,v_temp,tar_time])
                    self.world_x_list.append(x_temp)
                    self.world_y_list.append(y_temp)
                    self.world_theta_list.append(h_temp)
                i+=1
        self.traj_base_ego=self.traj
        self.world_x_list=np.array(self.world_x_list)
        self.world_y_list=np.array(self.world_y_list)
        self.world_theta_list=np.array(self.world_theta_list)
        return
    
    def world_pts_to_ego_pts(self,x,y,theta):
        x_convert_list=np.zeros((x.shape[0],1))
        y_convert_list=np.zeros((y.shape[0],1))
        theta_convert_list=np.zeros((theta.shape[0],1))
        ego_theta=self.ego_state.theta
        ego_x=self.ego_state.x
        ego_y=self.ego_state.y
        rotation_matrix = np.array([
        [np.cos(ego_theta), np.sin(ego_theta)],
        [-np.sin(ego_theta), np.cos(ego_theta)]])

        ego_unit_vec=[math.cos(ego_theta),math.sin(ego_theta)]
        ego2world_vec=[-ego_x,-ego_y]
        inner_prod=ego2world_vec[0]*ego_unit_vec[0]+ego2world_vec[1]*ego_unit_vec[1]
        cross_prod=ego2world_vec[1]*ego_unit_vec[0]-ego2world_vec[0]*ego_unit_vec[1]
        world_x_base_ego=inner_prod
        world_y_base_ego=cross_prod

        for i in range(x.shape[0]):
            world_point = np.array([x[i], y[i]])
            rotated_point = np.dot(rotation_matrix, world_point)
            local_point = rotated_point + np.array([world_x_base_ego,world_y_base_ego])
            # x_convert_list[i]=local_point[0]
            # y_convert_list[i]=local_point[1]
            # theta_convert_list[i]=self.normalize_angle(theta[i]-ego_theta)
            self.traj_base_ego[i][0]=local_point[0]
            self.traj_base_ego[i][1]=local_point[1]
            self.traj_base_ego[i][2]=self.normalize_angle(theta[i]-ego_theta)
        return

    def process(self):
        ilqr_proto_msg=self.GenerateIlqrProtoMsg()
        ilqr_params=self.LoadParams()
        proto_msg_result_string=pybind_test.RunIlqrController(ilqr_proto_msg.SerializeToString(),ilqr_params.SerializeToString())
        proto_msg_result=ilqr_controller_config_pb2.IlqrProtoMsg()
        proto_msg_result.ParseFromString(proto_msg_result_string)
        self.result_proto=proto_msg_result
        print("opt kappa[0]",proto_msg_result.ilqr_iteration_datas[len(proto_msg_result.ilqr_iteration_datas)-1].states[0].kappa)
        print("opt dkappa[0]",proto_msg_result.ilqr_iteration_datas[len(proto_msg_result.ilqr_iteration_datas)-1].controls[0].dkappa)
        print("opt kappa[1]",proto_msg_result.ilqr_iteration_datas[len(proto_msg_result.ilqr_iteration_datas)-1].states[1].kappa)
        print("opt dkappa[1]",proto_msg_result.ilqr_iteration_datas[len(proto_msg_result.ilqr_iteration_datas)-1].controls[1].dkappa)
        self.steer=self.Kappa2Steer(proto_msg_result.ilqr_iteration_datas[len(proto_msg_result.ilqr_iteration_datas)-1].states[1].kappa)
        #现在stm32的方向控制还是反的？
        self.steer=-self.steer
        atan_u_longi_normal=self.CalLongiControl()
        if atan_u_longi_normal==None:
            return False
        if atan_u_longi_normal>0:
            self.thr=atan_u_longi_normal
            self.brake=0
        else:
            self.thr=0
            self.brake=-atan_u_longi_normal

        return True
    

    def CalLongiControl(self):
        min_dis2=float('inf')
        match_idx=None
        for i,pt in enumerate(self.traj):
            temp_dis2=(pt[0]-self.ego_state.x)**2+(pt[1]-self.ego_state.y)**2
            if temp_dis2<min_dis2:
                min_dis2=temp_dis2
                match_idx=i
        if match_idx>=len(self.traj)-2:
            return None
        match_pt_s=self.traj[match_idx][3]
        match_pt_vec=[math.cos(self.traj[match_idx][2]),math.sin(self.traj[match_idx][2])]
        match_pt2ego_vec=[self.ego_state.x-self.traj[match_idx][0],self.ego_state.y-self.traj[match_idx][1]]
        inner_prod=match_pt_vec[0]*match_pt2ego_vec[0]+match_pt_vec[1]*match_pt2ego_vec[1]
        ego_s=match_pt_s+inner_prod

        num_of_timepiece=self.roi_time/self.dt
        tar_idx=int(num_of_timepiece)
        tar_s=None
        if tar_idx>=(len(self.traj)-1):
            tar_s=self.traj[-1][3]
        else:
            tar_s=self.traj[tar_idx][3]+(num_of_timepiece%1.0)*(self.traj[tar_idx+1][3]-self.traj[tar_idx][3])
        e_s=ego_s-tar_s

        tar_v=self.traj[match_idx][4]
        ego_v=self.ego_state.v
        e_v=ego_v-tar_v

        u_s=self.PID_control_s(e_s)
        u_v=self.PID_control_v(e_v)

        u_longi=self.dv_control_ratio*u_v+(1-self.dv_control_ratio)*u_s

        atan_u_longi_normal=math.atan(u_longi)*2/math.pi

        return atan_u_longi_normal


    def GenerateIlqrProtoMsg(self):
        ilqr_proto_msg=ilqr_controller_config_pb2.IlqrProtoMsg()
        traj_max_length=50
        for i in range(len(self.traj_base_ego)):
            if i>=traj_max_length:
                break
            traj_pt=ilqr_controller_config_pb2.TrajPoint()
            traj_pt.x=self.traj_base_ego[i][0]
            traj_pt.y=self.traj_base_ego[i][1]
            traj_pt.theta=self.traj_base_ego[i][2]
            traj_pt.v=self.traj_base_ego[i][3]
            ilqr_proto_msg.traj_points.add().CopyFrom(traj_pt)

        ego_state_proto=control_common_pb2.State()
        # ego_state_proto.x=self.ego_state.x
        # ego_state_proto.y=self.ego_state.y
        # ego_state_proto.theta=self.ego_state.theta
        ego_state_proto.x=0.0
        ego_state_proto.y=0.0
        ego_state_proto.theta=0.0
        ilqr_proto_msg.ego_state.CopyFrom(ego_state_proto)

        return ilqr_proto_msg
    
    def LoadParams(self):
        ilqr_params_path = '/home/fengmq/code/xiaoche_pi_new/smartcar/ilqr_control_pybind/proto_data/ilqr_params.pb.txt'
        ilqr_params = ilqr_controller_config_pb2.IlqrParams()
        with open(ilqr_params_path, 'r') as file:
            file_content = file.read()
            res = text_format.Parse(file_content, ilqr_params)
        print("success load parmas")
        return ilqr_params

    def update_ego_state(self, x ,y ,h, v, roi_time):
        # traj=np.array(self.traj)
        h=self.clockwise_0_360_to_pi_negative_pi(h)
        self.ego_state.update(x ,y ,h, v)
        self.roi_time=roi_time
        self.world_pts_to_ego_pts(self.world_x_list,self.world_y_list,self.world_theta_list)


    def normalize_angle(self, angle):
        """
        Normalize an angle to [-pi, pi].
        :param angle: (float)
        :return: (float) Angle in radian in [-pi, pi]
        """
        while angle > np.pi:
            angle -= 2.0 * np.pi

        while angle < -np.pi:
            angle += 2.0 * np.pi

        return angle

    def clockwise_0_360_to_pi_negative_pi(self,angle):
        angle=-angle
        angle=angle/180*math.pi
        while angle>=math.pi:
            angle-=2*math.pi
        while angle < - math.pi:
            angle += 2 * math.pi
        return angle
    
    def PID_control_s(self, e):
        self.Integral_record_ds.append(e)
        while(len(self.Integral_record_ds)>=self.KIntegralMaxSize):
            self.Integral_record_ds.pop(0)

        diffe_temp=(e-self.pre_ds)/0.1
        self.pre_ds=e

        thr_P=self.KPS*e
        thr_I=self.KIS*sum(self.Integral_record_ds)
        thr_D=self.KDS*diffe_temp

        u=thr_P+thr_I+thr_D
        u=-u
        return u
    
    def PID_control_v(self, e):
        self.Integral_record_dv.append(e)
        while(len(self.Integral_record_dv)>=self.KIntegralMaxSize):
            self.Integral_record_dv.pop(0)

        diffe_temp=(e-self.pre_dv)/0.1
        self.pre_dv=e

        thr_P=self.KPV*e
        thr_I=self.KIV*sum(self.Integral_record_dv)
        thr_D=self.KDV*diffe_temp

        u=thr_P+thr_I+thr_D
        u=-u
        return u
    
    def Kappa2Steer(self, kappa):
        #这一版忽略了beta角，后面要加上
        steer=math.atan((self.lf+self.lr)*kappa)
        if steer>self.KSteerLimit:
            steer=self.KSteerLimit
        elif steer<-self.KSteerLimit:
            steer=-self.KSteerLimit
        return steer